# uncompyle6 version 3.2.3
# Python bytecode 3.6 (3379)
# Decompiled from: Python 3.6.8 |Anaconda custom (64-bit)| (default, Feb 21 2019, 18:30:04) [MSC v.1916 64 bit (AMD64)]
# Embedded file name: scanners\wvs\wvs_scan_job.py
""" WVS scan job """
import os, logging, queue, json, time
from urllib.parse import urlsplit, urlunsplit, parse_qs, unquote
import zipfile, requests, requests.exceptions, shutil
from scanners.wvs.wvs_settings import WVSScanSettings, WVSScanSettingsError
from scanners.wvs.wvs_process import WVSProcess, WVSProcessError, WVSExitCode
from scanners.exceptions import *
from helpers.setup_types import is_setup_type_aws, is_setup_type_on_premise
from helpers.address import remove_query_parameters
from scanners.constants import text_data_identifier
from helpers.address import change_url_address
from settings_provider import settings
from scanners.constants import ScanJobStatus
from helpers.hashing.rand import get_random_lower_case_digits_string

_module_log = logging.getLogger("scan.wvs_scan_job")


class WVSScanningJob:
    """Class representing a scanning job with WVS"""

    CONTINUOUS_QUICK_PROFILE_NAME = "continuous_quick"
    RESULT_STATE_FILE_NAME = "result_state"
    V11_RESULT_FILE_NAME = "results.wvs"

    def __init__(self, job_description, app):
        """Initialize a scanning job"""
        _module_log.debug("initializing new job with id %s", job_description.job_id)
        self.status = ScanJobStatus.PREPARING
        self.app = app
        self.working_directory = os.path.join(
            app.working_directory,
            job_description.job_id + "_" + get_random_lower_case_digits_string(5),
        )
        self.description = job_description
        self.scan_settings = WVSScanSettings(job_description.main_target.address)
        self.fail_reason = ""
        self.scan_process = None
        self.debug = getattr(job_description, "debug", False) or settings.get(
            "debug_mode", False
        )
        self.is_marked_for_deletion = False
        self.archive = None
        self.scan_profile = ""
        self.initial_state_file_name = None
        self._WVSScanningJob__queue = queue.Queue()
        self.imported_files = []
        try:
            self.prepare()
        except Exception as e:
            self._fail("Job preparation failed")
            _module_log.exception("Job preparation failed with %s", e)

    def check_if_scan_is_still_possible(self):
        if self.status != ScanJobStatus.PREPARING:
            self.enqueue_long_operation(self._launch_scan)
            return
        if not self.description.worker_identifier:
            self.enqueue_long_operation(self._launch_scan)
            return
        time.sleep(1)
        from helpers.http_clients.status_api import ScanStatusClient, HTTP409, HTTP404

        scan_status_client = ScanStatusClient(
            endpoint=self.description.status_api_url,
            api_key=self.description.status_api_key,
        )
        try:
            result = scan_status_client.get_request(
                "/can_start_scan/%s" % (self.description.worker_identifier,)
            )
            if result["can_start"] == "true":
                self.enqueue_long_operation(self._launch_scan)
            else:
                if result["can_start"] == "wait":
                    self.enqueue_long_operation(self.check_if_scan_is_still_possible)
                else:
                    raise Exception(
                        "invalid can_start_scan response %s", result["can_start"]
                    )
        except HTTP409 as e:
            _module_log.warning(
                "stan stat check confirmation returned %s, aborting the scan", e
            )
            self.status = ScanJobStatus.ABORTED
            self.can_be_deleted()
        except HTTP404 as e:
            _module_log.warning(
                "stan stat check confirmation returned %s, aborting the scan", e
            )
            self.status = ScanJobStatus.ABORTED
            self.can_be_deleted()
        except Exception as e:
            _module_log.warning(
                "stan stat check confirmation returned %s, retrying the check", e
            )
            self.enqueue_long_operation(self.check_if_scan_is_still_possible)

    def enqueue_long_operation(self, method, *args, **kwargs):
        self._WVSScanningJob__queue.put((method, args, kwargs))

    def consume_enqueued_operations(self, consume_all=False):
        while 1:
            try:
                task = self._WVSScanningJob__queue.get(block=False)
            except queue.Empty:
                return

            if task is not None:
                try:
                    task[0](*task[1], **task[2])
                except Exception as error:
                    self._fail("task failed with %s" % (error,))

                self._WVSScanningJob__queue.task_done()
            if not consume_all:
                return

    def _fail(self, fail_reason):
        """Indicate job failure with a reason"""
        _module_log.debug(
            'Job %s failed because "%s"', self.description.job_id, fail_reason
        )
        self.fail_reason = fail_reason
        self.status = ScanJobStatus.FAILED

    def _add_dependent_file(self, value, file_name, ignore_errors=False):
        """Returns a local path to a required file and if necessary adds the
        URL to the download list. If the don't starts with a known protocol
        name like HTTP(s), FILE, TEXT_DATA or is a serializable dict we assume it is a local file with full path
        """
        if not value:
            return
        if isinstance(value, dict):
            profile_path = os.path.join(self.working_directory, file_name)
            (open(profile_path, "wt", encoding="utf-8")).write(json.dumps(value))
            return profile_path
        if value.startswith(text_data_identifier):
            profile_path = os.path.join(self.working_directory, file_name)
            (open(profile_path, "wt", encoding="utf-8")).write(
                value.replace(text_data_identifier, "", 1)
            )
            return profile_path
        _module_log.debug('Add dependent file "%s" as %s', str(value), file_name)
        scheme, _, path, query, _ = urlsplit(value)
        if scheme in ("http", "https"):
            q = parse_qs(query)
            ext = q.get("ext", [""])[0]
            secret = q.get("secret", [None])[0]
            value = remove_query_parameters(value, ("ext", "secret"))
            local_name = file_name + ext
            self.enqueue_long_operation(
                self._download_dependent_file,
                value,
                local_name,
                secret,
                ignore_errors=ignore_errors,
            )
            return os.path.join(self.working_directory, local_name)
        if scheme == "file":
            path = unquote(path[1:])
            ext = parse_qs(query).get("ext", [""])[0]
            if os.path.isfile(path):
                return path + ext
            if not ignore_errors:
                raise DependentFileMissing(path)
            else:
                return value

    def _download_dependent_file(self, url, file_name, secret, ignore_errors=False):
        """
        Download a required file and place it in the working_directory
        :param url:
        :param file_name:
        :param secret: if secret is present means our result is encrypted
        :return:
        """
        if self.status != ScanJobStatus.PREPARING:
            return
        if is_setup_type_on_premise():
            url = change_url_address(url, self.description.status_api_url)
        try:
            response = requests.get(url, stream=True, verify=False)
        except requests.exceptions.Timeout:
            if not ignore_errors:
                self._fail("Timeout while downloading %s" % url)
        except requests.exceptions.ConnectionError:
            if not ignore_errors:
                self._fail("Unable to connect to host for downloading %s" % url)
        except requests.exceptions.RequestException as error:
            if not ignore_errors:
                self._fail("Unable to download %s. Exception %s" % (url, error))
        else:
            if response.status_code == 200:
                destination_path = os.path.join(self.working_directory, file_name)
                with open(destination_path, "wb") as (output):
                    for chunk in response.iter_content(2048):
                        if self.status not in (
                            ScanJobStatus.PREPARING,
                            ScanJobStatus.EXECUTING,
                        ):
                            return
                        if chunk:
                            output.write(chunk)

                if secret:
                    from helpers.encryption.file_aes_cbc import (
                        decrypt_file,
                        InvalidData,
                        InvalidKey,
                    )

                    dec_path = os.path.join(self.working_directory, file_name + ".dec")
                    try:
                        decrypt_file(destination_path, dec_path, secret)
                    except InvalidData:
                        _module_log.warning(
                            "file %s does not appear to be encrypted", destination_path
                        )
                    except InvalidKey:
                        _module_log.error(
                            "file %s invalid decryption key", destination_path
                        )
                    else:
                        os.unlink(destination_path)
                        os.rename(dec_path, destination_path)

                else:
                    if not ignore_errors:
                        self._fail(
                            "Unable to download %s. Status returned %d."
                            % (url, response.status_code)
                        )

    def _launch_scan(self):
        """Launches a WVS instance with the prepared settings. This function should only be called
        from downloading dependent files or when prepare results in an empty download queue.
        _launch_scan assumes that self.scan_profile is filled (either with a downloaded profile path
        or with a known profile"""
        if self.status != ScanJobStatus.PREPARING:
            return
        self.status = ScanJobStatus.EXECUTING
        try:
            if self.scan_profile == WVSScanningJob.CONTINUOUS_QUICK_PROFILE_NAME:
                _module_log.debug(
                    "Launching quick part of a continuous scan on %s",
                    self.description.main_target.address,
                )
                self.scan_settings.patch_for_quick_continuous()
            if self.description.get("user_interaction", False):
                self.scan_settings.set_interactive_mode(True)
            demo_mode = self.description.get("mode") == "demo"
            savefile_name = self.RESULT_STATE_FILE_NAME
            if self.debug:
                if settings.get("features.engine_v11"):
                    savefile_name = self.V11_RESULT_FILE_NAME
                self.scan_process = WVSProcess(
                    wvs_directory=self.app.wvs_directory,
                    wvs_settings=self.scan_settings,
                    working_directory=self.working_directory,
                    profile=self.scan_profile,
                    status_api=self.description.status_api_url,
                    status_api_key=self.description.status_api_key,
                    log_level="debug",
                    savefile_name=savefile_name,
                    demo_mode=demo_mode,
                    initial_state_file_name=self.initial_state_file_name,
                )
                self.scan_process.debug = True
            else:
                if settings.get("features.engine_v11"):
                    savefile_name = None
                self.scan_process = WVSProcess(
                    wvs_directory=self.app.wvs_directory,
                    wvs_settings=self.scan_settings,
                    working_directory=self.working_directory,
                    profile=self.scan_profile,
                    status_api=self.description.status_api_url,
                    status_api_key=self.description.status_api_key,
                    savefile_name=savefile_name,
                    demo_mode=demo_mode,
                    initial_state_file_name=self.initial_state_file_name,
                )
            self.scan_process.start_scanning()
        except WVSProcessError as e:
            _module_log.exception(
                "Unable to start the scanner [%s]", e.__class__.__name__
            )
            self._fail("Unable to start the scanner [%s]" % (e.__class__.__name__,))

    def prepare(self):
        """Prepare for executing the job. Read WVS settings and download required files if
        necessary"""
        _module_log.debug('Preparing job "%s"', self.description.job_id)
        self.status = ScanJobStatus.PREPARING
        try:
            os.makedirs(self.working_directory)
        except FileExistsError:
            self._fail("Temporary folder already exists")
            return
        else:
            if hasattr(self.description, "scan_profile"):
                self.scan_profile = self._add_dependent_file(
                    self.description.scan_profile, "scan_profile"
                )
            initial_state_file_name = getattr(self.description, "initial_state", None)
            if initial_state_file_name:
                self.initial_state_file_name = self._add_dependent_file(
                    initial_state_file_name, "initial_state", ignore_errors=True
                )
            try:
                sequence_file = None
                if hasattr(self.description.main_target.configuration, "login"):
                    if (
                        self.description.main_target.configuration.login.kind
                        == "sequence"
                    ):
                        sequence_file = self._add_dependent_file(
                            self.description.main_target.get("login_sequence"),
                            "login.lsr",
                        )
                self.scan_settings.configure(
                    config=self.description.main_target.configuration,
                    login_sequence=sequence_file,
                    client_certificate=self._add_dependent_file(
                        self.description.main_target.get("client_certificate"), "cc"
                    ),
                )
                if hasattr(self.description, "bxss_user_id"):
                    if hasattr(self.description, "bxss_scan_id"):
                        self.scan_settings.setup_bxss(
                            self.description.bxss_user_id, self.description.bxss_scan_id
                        )
            except WVSScanSettingsError as error:
                self._fail(error.reason)
            except DependentFileError as error:
                self._fail(error.reason)

            try:
                if hasattr(self.description.main_target, "imports"):
                    for index, imported in enumerate(
                        self.description.main_target.imports
                    ):
                        filename = self._add_dependent_file(
                            imported, "import" + str(index)
                        )
                        if filename is not None:
                            self.imported_files.append(filename)
                            self.scan_settings.add_import(filename)

                if hasattr(self.description, "additional_targets"):
                    for index, additional_target in enumerate(
                        self.description.additional_targets
                    ):
                        self.scan_settings.add_allowed_host(
                            url=additional_target.address,
                            config=additional_target.configuration,
                            client_certificate=self._add_dependent_file(
                                additional_target.get("client_certificate"),
                                "cc" + str(index),
                            ),
                        )

            except WVSScanSettingsError as settings_error:
                self._fail(settings_error.reason)
            except DependentFileError as error:
                self._fail(error.reason)

            if self.status != ScanJobStatus.FAILED:
                if is_setup_type_aws():
                    self.enqueue_long_operation(self.check_if_scan_is_still_possible)
                else:
                    self.enqueue_long_operation(self._launch_scan)

    def clean_up(self):
        """Cleans up after job is finished"""
        self.delete_working_directory()

    def delete_working_directory(self):
        """Deletes the working_directory"""
        try:
            if os.path.isdir(self.working_directory):
                shutil.rmtree(self.working_directory)
        except PermissionError as error:
            _module_log.error("delete_working_directory failed with: %s", error)
        except OSError as error:
            _module_log.exception("delete_working_directory failed with: %s", error)

    def abort(self):
        """Abort the job"""
        if self.status == ScanJobStatus.PREPARING:
            self.status = ScanJobStatus.ABORTING
        else:
            if self.status == ScanJobStatus.PAUSING:
                pass
            else:
                if self.status == ScanJobStatus.EXECUTING:
                    if self.scan_process.is_alive():
                        self.scan_process.ctrl_break()
                        self.status = ScanJobStatus.ABORTING
                    else:
                        self.status = ScanJobStatus.ABORTED
                    return self.status

    def pause(self):
        """Pause job"""
        if self.status == ScanJobStatus.PREPARING:
            pass
        if self.status == ScanJobStatus.ABORTING:
            pass
        else:
            if self.status == ScanJobStatus.EXECUTING:
                if self.scan_process.is_alive():
                    self.scan_process.ctrl_c()
                    self.status = ScanJobStatus.PAUSING
                return self.status

    def update_status(self):
        """Checks for events with the scanning process and updates status accordingly"""
        just_finished = False
        if self.status == ScanJobStatus.EXECUTING:
            if not self.scan_process.is_alive():
                if self.scan_process.has_crashed():
                    if self.scan_process.exit_code == WVSExitCode.LICENSE_PROBLEM:
                        message = "License problem"
                    else:
                        if self.scan_process.exit_code == WVSExitCode.CRASH:
                            message = "Crash"
                        else:
                            if self.scan_process.exit_code == WVSExitCode.INIT_PROBLEM:
                                message = "Scan initialization problem"
                            else:
                                if (
                                    self.scan_process.exit_code
                                    == WVSExitCode.INVALID_COMMAND_LINE
                                ):
                                    message = "Invalid command line"
                                else:
                                    message = "Unhandled crash"
                                self._fail(
                                    dict(
                                        exit_code=self.scan_process.exit_code,
                                        message=message,
                                        stderr=self.scan_process.stderr,
                                    )
                                )
                else:
                    self.status = ScanJobStatus.FINISHED
                just_finished = True
        else:
            if self.status == ScanJobStatus.ABORTING:
                if not self.scan_process or not self.scan_process.is_alive():
                    self.status = ScanJobStatus.ABORTED
                    just_finished = True
                else:
                    if self.scan_process.ctrl_break_time_passed(240):
                        _module_log.error(
                            "Scanner does not respond to abort (ctrl-break), killing process"
                        )
                        self.scan_process.kill()
                        self.status = ScanJobStatus.ABORTED
                        just_finished = True
            else:
                if self.status == ScanJobStatus.PAUSING:
                    if not self.scan_process.is_alive():
                        self.status = ScanJobStatus.PAUSED
                        just_finished = True
                    else:
                        if self.scan_process.ctrl_c_time_passed(240):
                            _module_log.error(
                                "Scanner does not respond to pause (ctrl-c), killing process"
                            )
                            self.scan_process.kill()
                            self._fail("Process pausing failed")
                            just_finished = True
                    if just_finished:
                        _module_log.debug(
                            "job %s is just finished", self.description.job_id
                        )
                        self.create_archive()
                    return (self.status, just_finished)

    def delete_when_finished(self):
        """Mark this job to be deleted as soon as it reaches final states"""
        self.is_marked_for_deletion = True

    def can_be_deleted(self):
        """Returns True if the status is right for deleting the job"""
        return self.status in (
            ScanJobStatus.ABORTED,
            ScanJobStatus.FINISHED,
            ScanJobStatus.FAILED,
            ScanJobStatus.PAUSED,
        )

    def create_archive(self):
        if is_setup_type_aws():
            self.create_aws_archive()
        else:
            if is_setup_type_on_premise():
                self.create_on_premise_archive()

    def create_aws_archive(self):
        import subprocess
        from settings_provider import settings
        from helpers.aws.s3 import upload_file
        from helpers.hashing import sha256

        if self.status == ScanJobStatus.PAUSED or hasattr(
            self.description, "bxss_user_id"
        ):
            result_file_name = self.scan_process.get_filepath(
                self.scan_process.savefile_name
            )
            _module_log.debug("attempting to save: %s", result_file_name)
            if os.path.exists(result_file_name):
                pass
        try:
            upload_file(
                settings["storage.scan_state.s3.bucket"],
                "%s/%s/%s"
                % (
                    settings["storage.scan_state.s3.dir"],
                    self.app.scanning_app,
                    self.description.job_id,
                ),
                result_file_name,
            )
        except Exception as e:
            _module_log.exception("scan state db file save failed with %s", e)

        try:
            files = []
            files_to_archive = [
                self.scan_process.get_filepath(self.scan_process.logfile_name),
                self.scan_process.get_filepath(self.scan_process.savefile_name),
                self.scan_process.get_filepath(self.scan_process.settings_name),
                self.scan_process.http_log(),
                os.path.join(self.working_directory, "scan.profile"),
                os.path.join(self.working_directory, "login.lsr"),
            ]
            for filename in self.imported_files:
                files_to_archive.append(os.path.join(self.working_directory, filename))

            for filename in [
                f for f in os.listdir(self.working_directory) if f.endswith(".dmp")
            ]:
                files_to_archive.append(os.path.join(self.working_directory, filename))

            for filename in files_to_archive:
                if filename is not None:
                    if os.path.isfile(filename):
                        files.append(filename)

            if files:
                password = sha256(
                    settings["storage.scan_logs.s3.salt"] + self.description.job_id
                )
                archive_file_name = (
                    os.path.join(settings["storage.tmp"], self.description.job_id)
                    + ".7z"
                )
                args = [
                    "7z.exe",
                    "a",
                    "-Y",
                    "-p%s" % password,
                    "-mhe",
                    archive_file_name,
                ]
                for file in files:
                    args.append(file)

                process = subprocess.Popen(args, cwd=settings["storage.tmp"])
                while process.poll() is None:
                    time.sleep(1)

                upload_file(
                    settings["storage.scan_logs.s3.bucket"],
                    "%s/%s/%s.7z"
                    % (
                        settings["storage.scan_logs.s3.dir"],
                        self.description.job_id,
                        "wvs",
                    ),
                    archive_file_name,
                )
                os.unlink(archive_file_name)
        except Exception as e:
            _module_log.exception("scan archiving failed with %s", e)

    def create_on_premise_archive(self):
        _module_log.debug('Creating archive for "%s"', self.description.job_id)

        def handle_scan_state_file():
            if self.status == ScanJobStatus.PAUSED or hasattr(
                self.description, "bxss_user_id"
            ):
                result_file_name = self.scan_process.get_filepath(
                    self.scan_process.savefile_name
                )
                _module_log.debug("attempting to save: %s", result_file_name)
                scan_state_storage = settings.get("storage.scan_state_db")
                if not os.path.exists(result_file_name):
                    _module_log.warning("scan_state_db not found %s", result_file_name)
                else:
                    if scan_state_storage is None:
                        _module_log.warning("scan_state_storage is not defined")
                    else:
                        destination_path = os.path.join(
                            scan_state_storage, self.app.scanning_app
                        )
                try:
                    if not os.path.exists(destination_path):
                        os.makedirs(destination_path, exist_ok=True)
                    shutil.move(
                        self.scan_process.get_filepath(self.scan_process.savefile_name),
                        os.path.join(destination_path, self.description.job_id),
                    )
                except Exception as e:
                    _module_log.exception("scan status db file save failed with %s", e)

        def archive_scan_results_worker():
            try:
                if not self.app.scan_storage_directory is not None:
                    raise AssertionError
                zipfile_name = (
                    os.path.join(
                        self.app.scan_storage_directory, self.description.job_id
                    )
                    + ".zip"
                )
                files_to_archive = [
                    self.scan_process.get_filepath(self.scan_process.logfile_name),
                    self.scan_process.get_filepath(self.scan_process.settings_name),
                    self.scan_process.http_log(),
                    os.path.join(self.working_directory, "scan.profile"),
                    os.path.join(self.working_directory, "login.lsr"),
                ]
                for filename in self.imported_files:
                    files_to_archive.append(
                        os.path.join(self.working_directory, filename)
                    )

                for filename in [
                    f for f in os.listdir(self.working_directory) if f.endswith(".dmp")
                ]:
                    files_to_archive.append(
                        os.path.join(self.working_directory, filename)
                    )

                if self.debug:
                    files_to_archive.append(
                        self.scan_process.get_filepath(self.scan_process.savefile_name)
                    )
                with zipfile.ZipFile(zipfile_name, "w", zipfile.ZIP_DEFLATED) as (z):
                    archive_has_files = False
                    for filename in files_to_archive:
                        if filename is not None:
                            if os.path.isfile(filename):
                                z.write(filename, arcname=os.path.basename(filename))
                                archive_has_files = True

                if archive_has_files:
                    self.archive = zipfile_name
                else:
                    os.remove(zipfile_name)
                    self.archive = False
            except Exception as _e:
                self.archive = False
                _module_log.exception(_e)

        import gevent.monkey
        from threading import enumerate as enumerate_threads

        Thread = gevent.monkey.get_original("threading", "Thread")
        t = Thread(target=archive_scan_results_worker)
        t.start()
        while t in enumerate_threads():
            time.sleep(1)

        handle_scan_state_file()

    @property
    def extended_status(self):
        """Returns the extended status of the job"""
        if self.status not in (
            ScanJobStatus.ABORTED,
            ScanJobStatus.FINISHED,
            ScanJobStatus.FAILED,
            ScanJobStatus.PAUSED,
        ):
            return
        else:
            status = {}
            if self.fail_reason:
                status["fail_reason"] = self.fail_reason
            if self.archive:
                if os.path.isfile(self.archive):
                    status["attachments"] = [
                        {
                            "name": "output",
                            "url": urlunsplit(("file", "", self.archive, "", "")),
                        }
                    ]
            return status
